import json
import os

import cv2
from ultralytics import YOLO

from exchange_opt import *

slot_opt = SlotDetector(0.05, 150, 600)

# arrow
# blue 亮 200 暗 70 中
# red 亮 150 暗 50 中

slot_opt_box = SlotDetector(0.05, 150, 600)
# box
# blue 亮 200 暗 120 中
# red 亮 175 暗 100 中

image_path = "/home/champrin/Desktop/MV-CS016-10UC+DA1041860/给预备役/json_annotations (copy)/json_annotations/a"

# 获取文件夹A中所有图片文件路径
image_files = [os.path.join(image_path, f) for f in os.listdir(image_path) if
               f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

# 根据文件名进行排序
sorted_image_files = sorted(image_files, key=lambda x: os.path.basename(x))

# sorted_image_files = [r"/home/champrin/Desktop/MV-CS016-10UC+DA1041860/red/Video_20250113104416448/exchange_red_Video_20250113104416448_27.jpg"]


# 定义类别和对应的group_id
category_to_group_id = {
    "box": 0,
    "arrows_right": 1,
    "arrows_left": 2
}

# 打印排序后的图片文件路径
for file_path in sorted_image_files:
    print(file_path)

    rst = []

    classes = {'box': 0, 'arrow_right': 1, 'arrow_left': 2}


    def normalize(value, max_value):
        return value / max_value


    image = cv2.imread(file_path)
    image_height, image_width, _ = image.shape

    # cv2.line(image, (0, 0), (image_width, 00), (0, 0, 0), 10)
    # cv2.line(image, (image_width, 0), (image_width, image_height), (0, 0, 0), 10)
    # cv2.line(image, (image_width, image_height), (0, image_height), (0, 0, 0), 10)
    # cv2.line(image, (0, image_height), (0, 0), (0, 0, 0), 10)

    rst.append([
        float(0),
        float(0),
        float(image_width),
        float(image_height),
        int(2), float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
        float(1.0),
    ])

    # print(rst)

    slots = slot_opt.Detect(image, rst)
    slot_opt.DrawDebug(image)

    rst = []
    json_file_path = file_path.replace(".jpg", ".json")

    with open(json_file_path, 'r') as file:
        data = json.load(file)

    for shape in data["shapes"]:
        class_id = category_to_group_id.get(shape["label"])  # shape["label"] 类别名

        if class_id != category_to_group_id.get("box"):
            continue

        if shape["shape_type"] == "rectangle":
            # to xywh
            # find left/right corner
            x1, y1, x2, y2 = 0, 0, 0, 0
            if shape["points"][0][0] < shape["points"][1][0]:
                x1 = shape["points"][0][0]
                x2 = shape["points"][1][0]
            else:
                x1 = shape["points"][1][0]
                x2 = shape["points"][0][0]

            if shape["points"][0][1] < shape["points"][1][1]:
                y1 = shape["points"][0][1]
                y2 = shape["points"][1][1]
            else:
                y1 = shape["points"][1][1]
                y2 = shape["points"][0][1]

            width = abs(shape["points"][1][0] - shape["points"][0][0])
            height = abs(shape["points"][1][1] - shape["points"][0][1])
            # print(width, height)
            width = abs(x1 - x2)
            height = abs(y1 - y2)
            # print(width, height)

            keypoints = []
            for keypoint_label in [f"{shape['label']}_{i}" for i in range(1, 7)]:
                keypoint = next((keypoint for keypoint in data["shapes"] if keypoint["label"] == keypoint_label), None)
                if keypoint:
                    keypoint_x = keypoint["points"][0][0]
                    keypoint_y = keypoint["points"][0][1]
                    keypoint_visibility = 1.0  # Assuming all keypoints are visible 0：没有明显露出，不可见 1：被遮挡，不可见 2：可见
                    keypoints.extend([keypoint_x, keypoint_y, keypoint_visibility])
                else:
                    keypoints.extend([0, 0, 0])  # If keypoint not found, use 0 values

            rst.append([
                x1, y1, width, height,
                class_id, 1.0,
                keypoints[0], keypoints[1], keypoints[2],
                keypoints[3], keypoints[4], keypoints[5],
                keypoints[6], keypoints[7], keypoints[8],
                keypoints[9], keypoints[10], keypoints[11],
                keypoints[12], keypoints[13], keypoints[14],
                keypoints[15], keypoints[16], keypoints[17]
            ])

    # print(rst)
    if len(rst) == 0:
        print("rst == 0")
        continue

    slots_box = slot_opt_box.Detect(image, rst)
    slot_opt_box.DrawDebug(image)

    # print(slots)
    # print(slots_box)

    slots = slots + slots_box

    yolo_format_data = []
    for slot in slots:
        x, y, w, h = slot.rect

        center_x = x + w / 2
        center_y = y + h / 2
        width = w
        height = h

        # 归一化边界框坐标
        center_x = normalize(center_x, image_width)
        center_y = normalize(center_y, image_height)
        width = normalize(width, image_width)
        height = normalize(height, image_height)

        # 保留6位小数
        center_x = round(center_x, 6)
        center_y = round(center_y, 6)
        width = round(width, 6)
        height = round(height, 6)

        class_id = slot.type.value

        keypoints = []
        cnt = 0

        if class_id == 2:
            # 倒置顺序后的二维数组
            slot.opt_pts = slot.opt_pts[::-1]
        elif class_id == 1:
            # 将第一个元素移到最后一个位置
            slot.opt_pts = slot.opt_pts[1:] + slot.opt_pts[:1]

        for keypoint in slot.opt_pts:
            keypoint_x = keypoint[0]
            keypoint_y = keypoint[1]
            keypoint_visibility = 2  # Assuming all keypoints are visible 0：没有明显露出，不可见 1：被遮挡，不可见 2：可见
            keypoint_x = normalize(keypoint_x, image_width)
            keypoint_y = normalize(keypoint_y, image_height)
            # 保留6位小数
            keypoint_x = round(keypoint_x, 6)
            keypoint_y = round(keypoint_y, 6)
            cnt += 1
            keypoints.extend([keypoint_x, keypoint_y, keypoint_visibility])

        if class_id == 0:
            if cnt != 6:
                for i in range(6 - cnt):
                    keypoints.extend([20., 20., 0])
            try:
                yolo_format_data.append(f"{class_id} {center_x} {center_y} {width} {height} "
                                        f"{keypoints[0]} {keypoints[1]} {keypoints[2]} "
                                        f"{keypoints[3]} {keypoints[4]} {keypoints[5]} "
                                        f"{keypoints[9]} {keypoints[10]} {keypoints[11]} "
                                        f"{keypoints[15]} {keypoints[16]} {keypoints[17]} "
                                        f"{keypoints[6]} {keypoints[7]} {keypoints[8]} "
                                        f"{keypoints[12]} {keypoints[13]} {keypoints[14]}"
                                        )
            except IndexError as e:
                print(e)
        else:
            yolo_format_data.append(
                f"{class_id} {center_x} {center_y} {width} {height} " + " ".join(map(str, keypoints)))

    print(yolo_format_data)

    output_folder = os.path.join(os.path.dirname(file_path), "yolo_annotations")
    os.makedirs(output_folder, exist_ok=True)
    filename = os.path.basename(file_path)
    output_file_path = os.path.join(output_folder, filename.replace(".jpg", ".txt"))
    print(output_file_path)
    print(len(yolo_format_data))
    if len(yolo_format_data) != 0:

        with open(output_file_path, 'w') as output_file:
            for line in yolo_format_data:
                output_file.write(line + "\n")

        cv2.namedWindow("warp_img", 0)
        cv2.imshow("warp_img", slot_opt.warp_img)
        cv2.namedWindow("warp_roi", 0)
        cv2.imshow("warp_roi", slot_opt.warp_roi)
        cv2.namedWindow("warp_roib", 0)
        cv2.imshow("warp_roib", slot_opt.warp_roib)
    cv2.namedWindow("image", 0)
    cv2.imshow("image", image)
    cv2.namedWindow("binary_img_", 0)
    cv2.imshow("binary_img_", slot_opt.binary_img_)
    key = cv2.waitKey(0)
    if key == 27:
        break

cv2.destroyAllWindows()
